from pydantic import BaseModel, Field, model_validator
from typing import Optional, List
import winreg as reg
import os

# REG_SZ: 一个以        NULL        结尾的字符串。
# REG_EXPAND_SZ: 一个包含未展开环境变量的字符串，例如 % PATH %。
# REG_BINARY: 二进制数据。
# REG_DWORD: 32        位二进制数。
# REG_QWORD: 64        位二进制数。
# REG_MULTI_SZ: 一个由多个以        NULL
# 结尾的字符串组成的数组，最后跟两个        NULL        字符。
# 本包只处理str和int两种
def 获取注册表类型(value):
    """ 根据提供的值返回对应的注册表类型。
    :param value: 要存储的值
    :return: 注册表类型常量
    """
    if isinstance(value, str):
        return reg.REG_SZ
    elif isinstance(value, int):
        return reg.REG_DWORD
    elif isinstance(value, bytes):
        return reg.REG_BINARY
    else:
        raise TypeError("不支持的值类型")

# 此函数设置项和键通用
def 注册表设置项或项的子键(root, path, 项名=None, 默认值=None, 子键名=None, 子键值=None):
    """ 在指定的注册表路径下创建一个新的注册表项，并设置其默认值和字符串值。
    :param root: 注册表根键 (例如 reg.HKEY_CURRENT_USER)
    :param path: 注册表路径 (例如 "SOFTWARE\\Classes\\*\\shell")
    :param 项名: 要创建的注册表项名称 (可选)
    :param 默认值: 设置项的默认值 (可选)
    :param 子键名: 要添加的值的名称 (可选)
    :param 子键值: 要设置的值 (可选)
    :return: 如果成功创建项并添加值，返回 True；否则返回 False
    """
    if 子键名=="SubCommands":
        pass
    try:
        # 构建完整的路径
        full_path = os.path.join(path, 项名) if 项名 else path

        # 尝试打开现有的项，如果不存在则创建新项
        key = None
        try:
            key = reg.OpenKey(root, full_path, 0, reg.KEY_ALL_ACCESS)
            print(f"成功打开现有项: {full_path}")
        except FileNotFoundError:
            key = reg.CreateKey(root, full_path)
            print(f"成功创建新项: {full_path}")

        if key is None:
            raise Exception("无法打开或创建注册表项")

        with key:
            # 如果提供了默认值，则设置默认值
            if 默认值 is not None:
                reg_type = 获取注册表类型(默认值)
                reg.SetValueEx(key, "", 0, reg_type, 默认值)
                print(f"成功设置默认值: {默认值}")

            # 如果提供了子键名和子键值，则添加值
            if 子键名 is not None and 子键值 is not None:
                reg_type = 获取注册表类型(子键值)
                reg.SetValueEx(key, 子键名, 0, reg_type, 子键值)
                print(f"成功添加子键: {子键名} = {子键值}")
            elif  子键名 is not None and 子键值 is None:
                print(f"错误，子键值未给出")
                return False

    except PermissionError as e:
        print(f"权限不足，无法创建项或添加值: {e}")
        return False
    except Exception as e:
        print(f"创建项或添加值时发生错误: {e}")
        return False

    return True

# 内部用，此函数需要预打开项句柄，所以不对外
def _删除子键(key, 子键名):
    """ 删除指定注册表项中的子键（值）。
    :param key: 已打开的注册表项句柄
    :param 子键名: 要删除的子键名称
    :return: 如果成功删除子键，返回 True；否则返回 False
    """
    try:
        reg.DeleteValue(key, 子键名)
        print(f"成功删除子键: {子键名}")
        return True
    except WindowsError as e:
        print(f"无法删除子键: {子键名} ({e})")
        return False

# 删除键值，不影响项
def 注册表删除键值(root, path, 子键名, 项名=None):
    """ 在指定的注册表项中删除一个键值。
    :param root: 注册表根键 (例如 reg.HKEY_CURRENT_USER)
    :param path: 注册表路径 (例如 "SOFTWARE\\Classes\\*\\shell")
    :param 项名: 包含要删除键值的注册表项名称
    :param 子键名: 要删除的子键名称
    :return: 如果成功删除键值，返回 True；否则返回 False
    """
    full_path = os.path.join(path, 项名) if 项名 else path
    try:
        # 打开指定的注册表项
        with reg.OpenKey(root, full_path, 0, reg.KEY_ALL_ACCESS) as key:
            # 尝试删除指定的子键
            if not _删除子键(key, 子键名):
                return False
            print(f"成功从项 {full_path} 中删除子键: {子键名}")
            return True
    except FileNotFoundError:
        print(f"未找到项: {full_path}")
        return False
    except PermissionError as e:
        print(f"权限不足，无法删除子键: {e}")
        return False
    except Exception as e:
        print(f"删除子键时发生错误: {e}")
        return False

# 内部用
def _列出子项(root, path):
    """
    列出指定注册表项下的所有子项名称。

    :param root: 注册表根键 (例如 reg.HKEY_CURRENT_USER)
    :param path: 注册表路径 (例如 "SOFTWARE\\Classes\\*\\shell")
    :return: 子项名称列表
    """
    try:
        key = reg.OpenKey(root, path, 0, reg.KEY_READ)
        subkeys = []
        i = 0
        while True:
            try:
                subkey_name = reg.EnumKey(key, i)
                subkeys.append(subkey_name)
                i += 1
            except OSError:
                break
        reg.CloseKey(key)
        return subkeys
    except FileNotFoundError:
        print(f"未找到项: {path}")
        return []
# 将项删除，递归使出其下所有子项和键值
def 注册表删除项及其子项(root, path, 项名=None):
    """
    在指定的注册表路径下删除注册表项及其所有子项和子键。

    :param root: 注册表根键 (例如 reg.HKEY_CURRENT_USER)
    :param path: 注册表路径 (例如 "SOFTWARE\\Classes\\*\\shell")
    :param 项名: 要删除的注册表项名称 (可选)
    :return: 如果成功删除项及其子项，返回 True；否则返回 False
    """
    try:
        # 如果提供了项名，则构建完整的路径
        if 项名 is not None:
            full_path = os.path.join(path, 项名)
        else:
            full_path = path
        # 尝试打开现有的项
        try:
            key = reg.OpenKey(root, full_path, 0, reg.KEY_ALL_ACCESS)
            print(f"成功打开现有项: {full_path}")
        except FileNotFoundError:
            print(f"未找到项: {full_path}")
            return False

        # 递归删除所有子项
        subkeys = _列出子项(root, full_path)
        for subkey in subkeys:
            subkey_path = os.path.join(full_path, subkey)
            if not 注册表删除项及其子项(root, full_path, subkey):
                print(f"无法删除子项: {subkey_path}")
                reg.CloseKey(key)
                return False

        # 删除当前项中的所有子键（值）
        try:
            i = 0
            while True:
                try:
                    value_name = reg.EnumValue(key, i)[0]
                    if not _删除子键(key, value_name):
                        reg.CloseKey(key)
                        return False
                    i += 1
                except OSError:
                    break
        except WindowsError as e:
            print(f"无法枚举子键: {e}")

        # 关闭当前项
        reg.CloseKey(key)

        # 删除当前项
        try:
            reg.DeleteKey(root, full_path)
            print(f"成功删除项: {full_path}")
        except WindowsError as e:
            print(f"无法删除项: {e}")
            return False

    except PermissionError as e:
        print(f"权限不足，无法删除项或子键: {e}")
        return False
    except Exception as e:
        print(f"删除项或子键时发生错误: {e}")
        return False

    return True

# 右键菜单管理工具函数#################################################
class 右键菜单(BaseModel):
    菜单名称: Optional[str]=None
    菜单文字: str
    菜单图标: Optional[str]=None
    执行命令: Optional[str]=None
    下级菜单: Optional[List['右键菜单']] = None  # 默认为 None

def showitems(item: 右键菜单, tab=0):
    if item.下级菜单:
        print(tab * "\t", f"{item.菜单名称}【{item.菜单文字}】")
        for r in item.下级菜单:
            showitems(r, tab + 1)
    else:
        print(tab * "\t", item.菜单文字, f"   cmd:( {item.执行命令})")

index=0
def install_menu(menu:右键菜单,root=reg.HKEY_CURRENT_USER,path=r"SOFTWARE\Classes\*\shell"):
    global index
    if menu.菜单名称==None:
        menu.菜单名称=f"MyMenu_{index:03d}"
        index+=1
    注册表设置项或项的子键(root, path, menu.菜单名称, 子键名="MUIVerb", 子键值=menu.菜单文字)
    if menu.菜单图标!=None:
        注册表设置项或项的子键(root, path, menu.菜单名称, 子键名="icon", 子键值=menu.菜单图标)
    if menu.执行命令!=None:
        注册表设置项或项的子键(root, path, menu.菜单名称 + r"\\command", 默认值=menu.执行命令)
    if menu.下级菜单!=None:
        注册表设置项或项的子键(root, path, menu.菜单名称, 子键名="SubCommands",子键值="")
        注册表设置项或项的子键(root, path, menu.菜单名称 + r"\\shell")
        for rec in menu.下级菜单:
            install_menu(rec, root, os.path.join(path, menu.菜单名称, "shell"))

def uninstall_menu(菜单名列表:List[str],root=reg.HKEY_CURRENT_USER,path=r"SOFTWARE\Classes\*\shell"):
    for name in 菜单名列表:
        注册表删除项及其子项(root, path, 项名=name)

def setmaxfile(i=None,root=reg.HKEY_CURRENT_USER):
    # 解决Win10选择超过15个文件后右键菜单会变短
    # 注册表定位至HKEY_CURRENT_USER\Software\Microsoft\Windows\CurrentVersion\Explorer。
    # 右键点击Explorer项，选择    新建 - DWORD（32    位）值，命名为    MultipleInvokePromptMinimum。
    # 新建注册表值 右键点击MultipleInvokePromptMinimum选择修改，或者双击它，切换为
    # 十进制，填入需要批量操作的最大文件数量。
    if type(i)==int:
        注册表设置项或项的子键(root,r"Software\Microsoft\Windows\CurrentVersion\Explorer",子键名="MultipleInvokePromptMinimum",子键值=i)
    else:
        注册表删除键值(root,r"Software\Microsoft\Windows\CurrentVersion\Explorer",子键名="MultipleInvokePromptMinimum")


if __name__ == "__main__":
    # 定义菜单实例
    my菜单1 = 右键菜单(
        菜单名称="test001",  # 指定菜单名称
        菜单文字="我的单级菜单",
        菜单图标=r"C:\Users\E450\Downloads\idle.ico",  # 使用原始字符串
        执行命令=r"C:\Windows\system32\win32calc.exe"  # 使用原始字符串
    )
    my菜单2 = 右键菜单(
        菜单名称="test002",  # 指定菜单名称
        菜单文字="我的多级菜单",  # 这里没有提供菜单名称，会自动生成
        菜单图标=r"C:\Users\E450\Downloads\idle.ico",
        下级菜单=[
            右键菜单(菜单文字="二级菜单_1", 执行命令=r"C:\Windows\system32\win32calc.exe"),
            右键菜单(菜单文字="二级菜单_2", 执行命令=r"C:\Windows\system32\win32calc.exe"),
            右键菜单(
                菜单文字="二级菜单_3",
                下级菜单=[
                    右键菜单(菜单文字="三级菜单_1", 执行命令=r"C:\Windows\system32\win32calc.exe"),
                    右键菜单(菜单文字="三级菜单_2", 执行命令=r"C:\Windows\system32\win32calc.exe"),
                ]
            )
        ]
    )

    # 提示用户选择操作
    choice = input("请选择要执行的操作 (1: 安装菜单测试, 2: 移除菜单测试): ").strip()

    if choice == '1':
        print("正在执行安装菜单测试...")
        install_menu(my菜单1,path=r"SOFTWARE\Classes\*\shell")
        install_menu(my菜单1,path=r"SOFTWARE\Classes\Directory\shell")
        install_menu(my菜单2,path=r"SOFTWARE\Classes\*\shell")
        install_menu(my菜单2,path=r"SOFTWARE\Classes\Directory\shell")
        print("安装菜单测试完成。")
    elif choice == '2':
        print("正在执行移除菜单测试...")
        uninstall_menu(["test001"],path=r"SOFTWARE\Classes\*\shell")
        uninstall_menu(["test001"],path=r"SOFTWARE\Classes\Directory\shell")
        uninstall_menu(["test002"],path=r"SOFTWARE\Classes\*\shell")
        uninstall_menu(["test002"],path=r"SOFTWARE\Classes\Directory\shell")
        print("移除菜单测试完成。")
    else:
        print("无效的选择，请输入 1 或 2。")
